import torch
import torch.nn as nn

# 超参数
input_size = 10   # 输入特征维度
hidden_size = 20  # 隐藏层维度
num_layers = 2    # GRU层数
batch_size = 3
seq_len = 5       # 序列长度

# 创建GRU模型 (batch_first=True表示输入为[batch, seq, feature])
gru = nn.GRU(
    input_size=input_size,
    hidden_size=hidden_size,
    num_layers=num_layers,
    batch_first=True,   # 批次维度在前
    bidirectional=False # 单向GRU
)

# 生成测试数据
inputs = torch.randn(batch_size, seq_len, input_size)  # [3, 5, 10]
h0 = torch.zeros(num_layers, batch_size, hidden_size)  # 初始隐藏状态 [2, 3, 20]

# 前向传播
output, hn = gru(inputs, h0)

print("输出维度:", output.shape)    # [3, 5, 20] (batch, seq_len, hidden_size)
print("最终隐藏状态:", hn.shape)     # [2, 3, 20] (num_layers, batch, hidden_size)